from .image_convert_base import ConvertBase
import cv2 as cv
import numpy as np

class FlipConvert(ConvertBase):
    
    #flip_type 1 ： 水平， 0 ：垂直， -1 ：对角线
    def __init__(self, flip_type = 1, use_rate = 0.5, is_rand = True):
        super().__init__(use_rate)
        self.flip_type = flip_type
        self.is_rand = is_rand
        
    
    def __get_flip_type(self):
        f_p = np.random.randint(-1, 2)   
        # print('flip_type : ', f_p)
        return f_p
    
    
    def convert(self, img, boxes, points):
        if self.is_rand:
            self.flip_type = self.__get_flip_type()
            
        img = cv.flip(img, self.flip_type)    
        h, w, _ = img.shape
        if boxes is not None:
            boxes = boxes.astype(np.float32)
            #水平
            if self.flip_type == 1:
                boxes[:, [0, 2]] = w - boxes[:, [2, 0]]
            elif self.flip_type == 0:
                boxes[:, [1, 3]] = h - boxes[:, [3, 1]]
            elif self.flip_type == -1:
                boxes[:, [0, 2]] = w - boxes[:, [2, 0]]
                boxes[:, [1, 3]] = h - boxes[:, [3, 1]]
        if points is not None:
            #水平
            if self.flip_type == 1:
                points[:, 0] = w - points[:, 0]
            elif self.flip_type == 0:
                points[:, 1] = h - points[:, 1]
            elif self.flip_type == -1:
                points[:, 0] = w - points[:, 0]
                points[:, 1] = h - points[:, 1]
        
        return img, boxes, points